/**
 * @file main.cpp
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-10-26
 *
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */

#include "imageDecoder.h"
#include "imageEncoder.h"
#include "ipeParamYolov3.h"
#include "opencvWindow.hpp"
#include "osd_plugin.h"
#include "util.hpp"
#include "yolox_post_process.h"
#include <dlfcn.h>
#include <lyn_api.h>
#include <nlohmann/json.hpp>
#include <string>
#include "CLI/CLI11.hpp"

using namespace std;

using YOLOX_POST_FUNC = int (*)(YoloxPostProcessInfo_t *);

static void hb_dram_objects(cv::Mat &img, lynBoxesInfo *objects)
{
  cv::Rect_<float> bbox;
  for (size_t i = 0; i < objects->boxesNum; ++i)
  {
    const lynBox &object = objects->boxes[i];
    bbox.x = object.xmin;
    bbox.y = object.ymin;
    bbox.width = object.xmax - object.xmin;
    bbox.height = object.ymax - object.ymin;
    std::string labelStr(object.label);

    cv::rectangle(img, bbox, cv::Scalar(255, 255, 0), 2);
    cv::Point textPosition(bbox.x, bbox.y - 10);
    cv::putText(img, labelStr, textPosition, cv::FONT_HERSHEY_SIMPLEX, 0.9,
                cv::Scalar(255, 255, 0), 2);
  }
}

void ShowImageWithCV(const std::string &strTitle, uint8_t *pImgData, int iWidth,
                     int iHeight, int iMiliSec)
{
  // cv::Mat imgSrc = cv::Mat(iWidth, iHeight, CV_8UC1, pImgData);
  cv::Mat imgSrc(iHeight * 3 / 2, iWidth, CV_8UC1, pImgData);

  cv::Mat yuvImg = imgSrc.clone();
  cv::Mat rgbimg(iHeight, iWidth, CV_8UC3);
  cvtColor(yuvImg, rgbimg, cv::COLOR_YUV2BGR_NV12);
  cv::imshow(strTitle, rgbimg);

  if (iMiliSec > 0)
  {
    cv::waitKey(iMiliSec);
  }
}

int main(int argc, char *argv[])
{
  CLI::App app;
  std::string inputFile;
  std::string outputFile;
  int deviceId;
  int encodeFlag;
  std::string postPlugin;
  std::string yolov3Model;
  std::string labelFile;
  std::string osdPlugin;

  app.add_option("input_image", inputFile, "input jpeg file path")->default_val("../data/yolov3_detect_image_with_opencv/test.jpeg");
  app.add_option("output_image", outputFile, "output jpeg file path")->default_val("../output.jpeg");
  app.add_option("device_id", deviceId, "lynxi device id")->default_val(0);
  app.add_option("encode_flag", encodeFlag, "0: show test picture directly with OpenCV, 1: save test result to be a picture file")->default_val(1);
  app.add_option("-p,--post_plugin", postPlugin, "yolov3 post process plugin path")->default_val("../common/plugin/postprocess/lib/libYolov3PostPlugin.so");
  app.add_option("-m,--yolov3_model", yolov3Model, "yolov3 model path")->default_val("../model/yolov3_b1_h608_w608_c3_iuint8_ofloat16_0000_quant/Net_0");
  app.add_option("-l,--label_file", labelFile, "yolov3 label path")->default_val("../common/plugin/postprocess/yolox/label/label.v3.json");
  app.add_option("-o,--osd_plugin", osdPlugin, "osd plugin path")->default_val("../common/plugin/osd/lib/libOsdPlugin.so");

  // 解析命令行参数
  CLI11_PARSE(app, argc, argv);

  if (!IsFileExist(inputFile))
  {
    cout << "error : input jpeg file is not exsit." << endl;
    return 0;
  }

  if (!IsValidDevice(deviceId))
  {
    cout << "error : device id is invalid." << endl;
    return 0;
  }

  if (!IsFileExist(postPlugin))
  {
    cout << "error : yolov3 post process plugin is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(yolov3Model))
  {
    cout << "error : yolov3 model is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(labelFile))
  {
    cout << "error : label file is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(osdPlugin))
  {
    cout << "error : osd plugin is not exsit." << endl;
    return 0;
  }

  ifstream fJson(labelFile);
  stringstream ss;
  ss << fJson.rdbuf();
  nlohmann::json js = nlohmann::json::parse(ss.str());

  int labelCount = js.size();
  if (labelCount == 0)
  {
    cout << "error : label file is not invalid." << endl;
    return 0;
  }

  LabelList labelList;
  labelList.labelNum = labelCount;
  labelList.labels = new LabelInfo[labelCount];
  memset(labelList.labels, 0, sizeof(LabelInfo) * labelCount);
  string labelName;
  for (int i = 0; i < labelCount; ++i)
  {
    labelName = js[i]["desc"].get<string>();
    strncpy(labelList.labels[i].name, labelName.c_str(),
            sizeof(labelList.labels[i].name) - 1);
  }

  // 2. 创建 context 与 stream
  lynContext_t ctx = nullptr;
  CHECK_ERR(lynCreateContext(&ctx, deviceId));
  CHECK_ERR(lynRegisterErrorHandler(StreamErrorHandler, nullptr));
  lynStream_t stream = nullptr;
  CHECK_ERR(lynCreateStream(&stream));

  // 3. 初始化解码图片类
  ImageDecoder imageDec;
  imageDec.Init(inputFile);
  lynImageInfo_t imgInfo;
  imageDec.GetImageDecInfo(&imgInfo);

  // 4. 同步解码图片到 Device 侧
  lynCodecBuf_t decodeImg;
  decodeImg.size = imgInfo.output.predictBufSize;
  CHECK_ERR(lynMalloc((void **)&decodeImg.data, decodeImg.size));
  imageDec.DecodeImageToDevice(stream, decodeImg.data);

  // 5. 加载模型
  ModelInfo yolov3ModelInfo;
  yolov3ModelInfo.LoadModelByPath(yolov3Model.c_str());
  lynPixelFormat_t pixelFormat = LYN_PIX_FMT_NV12;

  // 6. IPE 处理
  uint8_t *pIpeBufOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pIpeBufOut, yolov3ModelInfo.inputSize));
  IpeParamYolov3 ipeYolov3(yolov3ModelInfo.width, yolov3ModelInfo.height);
  ipeYolov3.SetImgInfo(imgInfo.output.width, imgInfo.output.height,
                       pixelFormat);
  ipeYolov3.CalcParam(stream, decodeImg.data, pIpeBufOut);

  // 7. 调用 APU 推理接口
  uint8_t *pDevApuOutData = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDevApuOutData, yolov3ModelInfo.outputSize));
  CHECK_ERR(lynExecuteModelAsync(stream, yolov3ModelInfo.model, pIpeBufOut,
                                 pDevApuOutData, yolov3ModelInfo.batchSize));

  // 8. 加载post Plugin, 进行模型后处理
  lynBoxesInfo hostBoxesInfo;
  YoloxPostProcessInfo_t post_info;
  memset(&post_info, 0, sizeof(YoloxPostProcessInfo_t));

  post_info.is_pad_resize = 1;
  post_info.score_threshold = 0.25;
  post_info.nms_threshold = 0.45;
  post_info.nms_top_k = 500;
  post_info.anchorSize = 3;
  post_info.width = yolov3ModelInfo.width;
  post_info.height = yolov3ModelInfo.height;
  post_info.ori_width = imgInfo.output.width;
  post_info.ori_height = imgInfo.output.height;
  post_info.class_num = labelCount;

  void *handle = dlopen(postPlugin.c_str(), RTLD_NOW | RTLD_GLOBAL);
  if (handle == nullptr)
  {
    cout << "error : post plugin load fail." << endl;
    return 0;
  }
  YOLOX_POST_FUNC yoloxPostFunc =
      (YOLOX_POST_FUNC)dlsym(handle, "lynYoloxPostProcess");
  if (yoloxPostFunc == nullptr)
  {
    cout << "error : post plugin has no lynYoloxPostProcess function." << endl;
    dlclose(handle);
    return 0;
  }

  lynSynchronizeStream(stream);
  uint8_t *pHostApuOutData = new uint8_t[yolov3ModelInfo.outputSize];
  CHECK_ERR(lynMemcpy(pHostApuOutData, pDevApuOutData,
                      yolov3ModelInfo.outputSize, ServerToClient));

  post_info.output_tensor = pHostApuOutData;
  post_info.labelList = labelList;

  lynBoxesInfo *boxesInfo = new lynBoxesInfo;
  post_info.boxesInfo = boxesInfo;

  yoloxPostFunc(&post_info);

  uint8_t *pHostDecBuf = new uint8_t[decodeImg.size];
  CHECK_ERR(lynMemcpyAsync(stream, pHostDecBuf, decodeImg.data, decodeImg.size,
                           ServerToClient));
  lynSynchronizeStream(stream);
  cv::Mat yuvImg(imgInfo.output.height * 3 / 2, imgInfo.output.width, CV_8UC1,
                 pHostDecBuf);

  hb_dram_objects(yuvImg, boxesInfo);

  delete boxesInfo;
  delete[] pHostApuOutData;
  dlclose(handle);

  // 等待 plugin 处理完，释放内存
  CHECK_ERR(lynFree(pDevApuOutData));
  CHECK_ERR(lynFree(pIpeBufOut));

  if (encodeFlag)
  {
    // 9. 进行图片编码
    CHECK_ERR(lynMemcpyAsync(stream, decodeImg.data, pHostDecBuf,
                             decodeImg.size, ClientToServer));
    ImageEncoder encoder;
    encoder.EncodeImageAndSave(stream, outputFile, decodeImg, imgInfo);
  }
  else
  {
    ShowImageWithCV("Yolov3Image", pHostDecBuf, imgInfo.output.width,
                    imgInfo.output.height, 10 * 1000);
  }

  delete[] pHostDecBuf;

  // 10. 销毁资源
  CHECK_ERR(lynDestroyStream(stream));
  CHECK_ERR(lynFree(decodeImg.data));
  yolov3ModelInfo.UnLoadModel();

  if (ctx)
  {
    CHECK_ERR(lynDestroyContext(ctx));
  }

  delete[] labelList.labels;

  return 0;
}
